package com.mxsoft;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

public class PerformanceTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        byte[] transformed = null;
        System.out.println("Transforming " + className);
        ClassPool pool = ClassPool.getDefault();
        CtClass ctClass = null;

        try {
            ctClass = pool.makeClass(new java.io.ByteArrayInputStream(classfileBuffer));

            if (ctClass.isInterface() == false) {
                CtBehavior[] behaviors = ctClass.getDeclaredBehaviors();
                for (CtBehavior ctBehavior : behaviors) {
                    if (ctBehavior.isEmpty()) {
                        doBehavior(ctBehavior);
                    }
                }

                transformed = ctClass.toBytecode();
            }
        } catch (CannotCompileException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (ctClass != null) {
                ctClass.detach();
            }
        }
        return transformed;
    }

    private void doBehavior(CtBehavior behavior) throws CannotCompileException {
        behavior.instrument(new ExprEditor() {
            @Override
            public void edit(MethodCall m) throws CannotCompileException {
                m.replace("{ long stime = System.nanoTime(); $_ = $proceed($$); System.out.println(" + m.getClassName()+"."+m.getMethodName() + ":"+"(System.nanoTime()-stime));}");
            }
        });
    }
}
